import torch
import torch.nn as nn
from torch import optim


class HAR_CNN(nn.Module):
    def __init__(self, input_dim, output_dim):
        """
        参数:
        - input_dim: 输入特征的维度 (561)
        - num_classes: 分类的类别数
        """
        super(HAR_CNN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv1d(1, 32, kernel_size=3, stride=1, padding=1),  # 卷积层，输入通道1，输出通道32
            nn.ReLU(),  # 激活函数 ReLU
            nn.MaxPool1d(kernel_size=2, stride=2),  # 最大池化，减少特征维度
            nn.Conv1d(32, 64, kernel_size=3, stride=1, padding=1),  # 第二个卷积层
            nn.ReLU(),  # 激活函数 ReLU
            nn.MaxPool1d(kernel_size=2, stride=2),  # 最大池化
        )
        self.fc = nn.Sequential(
            nn.Flatten(),  # 展平
            nn.Linear(64 * (input_dim // 4), 128),  # 全连接层
            nn.ReLU(),  # 激活函数 ReLU
            nn.Dropout(0.5),  # Dropout 防止过拟合
            nn.Linear(128, output_dim)  # 输出层，分类任务
        )

    def forward(self, x):
        x = self.conv(x)  # 卷积部分
        x = self.fc(x)  # 全连接部分
        return x
